import torch
import torch.nn as nn
from torchvision import models

__all__ = ["backbone_factory", "YoloV1Backbone", "Resnet18", "Resnet50"]


class GoogleNet(nn.Module):
    """
    Note: The model is GoogleNet-like model which mentioned in You only look once (https://www.cv-foundation.org/openaccess/content_cvpr_2016/papers/Redmon_You_Only_Look_CVPR_2016_paper.pdf). It needs to be pretrained with imagenet.
    """

    def __init__(self, class_num) -> None:
        super().__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3),
            nn.LeakyReLU(0.1),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(64, 192, kernel_size=3, padding=1),
            nn.LeakyReLU(0.1),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(192, 128, kernel_size=1),
            nn.LeakyReLU(0.1),
            nn.Conv2d(128, 256, kernel_size=3, padding=1),
            nn.LeakyReLU(0.1),
            nn.Conv2d(256, 256, kernel_size=1),
            nn.LeakyReLU(0.1),
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.LeakyReLU(0.1),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        self.conv4 = nn.Sequential(
            nn.Conv2d(512, 256, kernel_size=1),
            nn.LeakyReLU(0.1),
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.LeakyReLU(0.1),
            nn.Conv2d(512, 256, kernel_size=1),
            nn.LeakyReLU(0.1),
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.LeakyReLU(0.1),
            nn.Conv2d(512, 256, kernel_size=1),
            nn.LeakyReLU(0.1),
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.LeakyReLU(0.1),
            nn.Conv2d(512, 256, kernel_size=1),
            nn.LeakyReLU(0.1),
            nn.Conv2d(256, 512, kernel_size=3, padding=1),
            nn.LeakyReLU(0.1),
            nn.Conv2d(512, 512, kernel_size=1),
            nn.LeakyReLU(0.1),
            nn.Conv2d(512, 1024, kernel_size=3, padding=1),
            nn.LeakyReLU(0.1),
            nn.MaxPool2d(kernel_size=2, stride=2),
        )
        self.conv5 = nn.Sequential(
            nn.Conv2d(1024, 512, kernel_size=1),
            nn.LeakyReLU(0.1),
            nn.Conv2d(512, 1024, kernel_size=3, padding=1),
            nn.LeakyReLU(0.1),
            nn.Conv2d(1024, 512, kernel_size=1),
            nn.LeakyReLU(0.1),
            nn.Conv2d(512, 1024, kernel_size=3, padding=1),
            nn.LeakyReLU(0.1),
        )
        self.classifier = nn.Sequential(
            nn.AvgPool2d(kernel_size=1), nn.Flatten(), nn.Linear(1024, class_num)
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        x = self.conv1(x)
        x = self.conv2(x)
        x = self.conv3(x)
        x = self.conv4(x)
        x = self.conv5(x)
        x = self.classifier(x)
        return x


class Backbone(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.backbone: nn.Module

    @property
    def feature_channels(self) -> int:
        input_tensor = torch.rand(1, 3, 448, 448)
        return self.forward(input_tensor).shape[1]

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.backbone(x)


class YoloV1Backbone(Backbone):
    """The backbone of yolo mentioned in You only look once (https://www.cv-foundation.org/openaccess/content_cvpr_2016/papers/Redmon_You_Only_Look_CVPR_2016_paper.pdf). It needs to be pretrained with imagenet."""

    def __init__(self) -> None:
        super().__init__()
        net = GoogleNet(class_num=2000)
        net.load_state_dict(torch.load("models/googlenet.pth"))
        self.backbone = nn.Sequential(
            net.conv1, net.conv2, net.conv3, net.conv4, net.conv5
        )


class YoloV2BackBone(Backbone):
    def __init__(self) -> None:
        super().__init__()
        # TODO: yolov2 backbone


class Resnet18(Backbone):
    def __init__(self) -> None:
        super().__init__()
        resnet18 = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)
        self.backbone = nn.Sequential(*list(resnet18.children())[:-2])


class Resnet50(Backbone):
    def __init__(self) -> None:
        super().__init__()
        resnet50 = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
        self.backbone = nn.Sequential(*list(resnet50.children())[:-2])


def backbone_factory(backbone: str):
    if backbone == "yolov1backbone":
        return YoloV1Backbone()
    elif backbone == "yolov2backbone":
        return YoloV2BackBone()
    elif backbone == "resnet18":
        return Resnet18()
    elif backbone == "resnet50":
        return Resnet50()
    else:
        raise NotImplementedError(
            "['yolov1backbone', 'yolov2backbone', 'resnet18', 'resnet50'] is supported only."
        )
